#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
迁移与回填 annotations.anchor_json 的脚本
- 备份数据库
- 为无 version 的旧记录补充 version=1，并确保 JSON 至少包含 method 与 quote 字段
- 可选：尝试为旧记录补充 context_before/context_after/occurrenceIndex 的近似值（启用需传参 --enhance）
运行:
  python migrate_annotations_anchor.py [--db data/app.sqlite或你的db路径] [--enhance]
"""
import os
import json
import shutil
import argparse
import datetime
import sqlite3

DEFAULT_DB = None  # 将在运行时从 db_migration.DATABASE 或常见路径探测

def detect_db_path():
    # 常见路径探测
    candidates = []
    try:
        from db_migration import DATABASE as DB_CONST  # type: ignore
        candidates.append(DB_CONST)
    except Exception:
        pass
    candidates.extend([
        os.path.join('data', 'app.sqlite'),
        os.path.join('data', 'database.sqlite'),
        os.path.join('Productions_DB', 'data', 'app.sqlite'),
    ])
    for p in candidates:
        if p and os.path.exists(p):
            return p
    # 回退：取 data 目录下第一个 .sqlite
    if os.path.isdir('data'):
        for name in os.listdir('data'):
            if name.lower().endswith('.sqlite') or name.lower().endswith('.db'):
                return os.path.join('data', name)
    return None

def backup_db(db_path):
    ts = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    dst_dir = 'data'
    os.makedirs(dst_dir, exist_ok=True)
    dst = os.path.join(dst_dir, f'db_backup_{ts}.sqlite')
    shutil.copy2(db_path, dst)
    return dst

def ensure_table_columns(conn):
    # 确认 annotations 表存在 anchor_json 列
    cur = conn.execute("PRAGMA table_info(annotations)")
    cols = [r[1] for r in cur.fetchall()]
    if 'anchor_json' not in cols:
        raise RuntimeError("annotations 表缺少 anchor_json 列，请先执行主数据库迁移（init_database）。")
    return cols

def fetch_content_map(conn):
    # 构建 content_id -> 合并文本（用于近似 context 回填；非强依赖）
    # 从 content.content 或 extracted_text 或 raw_content 汇总
    m = {}
    try:
        cur = conn.execute("""
            SELECT id, COALESCE(content, extracted_text, raw_content, '') AS txt
            FROM content
        """)
        for cid, txt in cur.fetchall():
            m[str(cid)] = txt or ''
    except Exception:
        pass
    return m

def limit_str(s, n=200):
    try:
        s = str(s or '')
        return s[:n]
    except Exception:
        return ''

def enhance_context(quote, full_text):
    # 尝试在全文中寻找 quote 的最后一次出现，构造上下文；若找不到返回空
    if not quote or not full_text:
        return None
    pos = full_text.rfind(quote)
    if pos < 0:
        pos = full_text.find(quote)
        if pos < 0:
            return None
    before = full_text[max(0, pos-20):pos]
    after = full_text[pos+len(quote):pos+len(quote)+20]
    return {
        "context_before": limit_str(before, 200),
        "context_after": limit_str(after, 200),
        "occurrenceIndex": pos
    }

def migrate(db_path, enhance=False):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    try:
        ensure_table_columns(conn)
        # 备份
        backup = backup_db(db_path)
        print(f"[备份完成] -> {backup}")

        content_map = fetch_content_map(conn) if enhance else {}

        cur = conn.execute("SELECT id, content_id, anchor_json FROM annotations")
        rows = cur.fetchall()
        print(f"[扫描] annotations 记录数: {len(rows)}")

        updated = 0
        for r in rows:
            aid = r['id']
            cid = r['content_id']
            raw = r['anchor_json'] or ''
            try:
                anchor = json.loads(raw) if raw else {}
            except Exception:
                # 非 JSON：尽力包裹
                anchor = {"method": "text-quote", "quote": limit_str(str(raw), 200)}

            # 若已经有 version >= 2，跳过（视为已是新结构）
            if isinstance(anchor, dict) and int(anchor.get('version', 0)) >= 2:
                continue

            changed = False
            # 确保 method/quote 存在
            if 'method' not in anchor:
                anchor['method'] = 'text-quote'; changed = True
            if 'quote' not in anchor:
                # 没 quote 基本无法定位，保留为空串
                anchor['quote'] = limit_str('', 200); changed = True
            else:
                # 裁剪长度
                q = limit_str(anchor.get('quote'), 200)
                if q != anchor.get('quote'):
                    anchor['quote'] = q; changed = True

            # 可选：为旧数据补充近似上下文
            if enhance and anchor.get('quote'):
                txt = content_map.get(str(cid), '')
                ctx = enhance_context(anchor['quote'], txt)
                if ctx:
                    if not anchor.get('context_before'):
                        anchor['context_before'] = ctx['context_before']; changed = True
                    if not anchor.get('context_after'):
                        anchor['context_after'] = ctx['context_after']; changed = True
                    if 'occurrenceIndex' not in anchor:
                        anchor['occurrenceIndex'] = ctx['occurrenceIndex']; changed = True

            # 标注版本
            if 'version' not in anchor or int(anchor.get('version', 0)) < 1:
                anchor['version'] = 1; changed = True

            if changed:
                conn.execute(
                    "UPDATE annotations SET anchor_json = ? , updated_at = CURRENT_TIMESTAMP WHERE id = ?",
                    (json.dumps(anchor, ensure_ascii=False), aid)
                )
                updated += 1

        conn.commit()
        print(f"[完成] 已更新 {updated} 条记录，其余保持不变。")
        print("说明：version=1 表示旧结构；新建/更新的批注会写入 version=2 的增强锚点。")
    finally:
        conn.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="迁移与回填 annotations.anchor_json")
    parser.add_argument("--db", type=str, default=None, help="数据库文件路径，默认自动探测")
    parser.add_argument("--enhance", action="store_true", help="为旧数据近似补充上下文和 occurrenceIndex")
    args = parser.parse_args()

    db_path = args.db or detect_db_path()
    if not db_path or not os.path.exists(db_path):
        print("未找到数据库文件。请使用 --db 指定路径，例如: --db data/app.sqlite")
        raise SystemExit(1)

    print(f"使用数据库: {db_path}")
    migrate(db_path, enhance=args.enhance)